use std::sync::Arc;

use async_trait::async_trait;
use pingora::server::Server;
use pingora::{prelude::*, upstreams};

fn main() {
    let mut pingora_server = Server::new(None).unwrap();
    pingora_server.bootstrap();

    let upstreams = LoadBalancer::try_from_iter(["127.0.0.1:8080", "127.0.0.2:8080"]).unwrap();
    let mut lb = http_proxy_service(&pingora_server.configuration, LB(Arc::new(upstreams)));
    lb.add_tcp("0.0.0.0:6188");
    pingora_server.add_service(lb);

    pingora_server.run_forever();
}

struct LB(Arc<LoadBalancer<RoundRobin>>);

#[async_trait]
impl ProxyHttp for LB {
    type CTX = ();
    fn new_ctx(&self) -> () {
        ()
    }

    async fn upstream_peer(&self, _session: &mut Session, _ctx: &mut ()) -> Result<Box<HttpPeer>> {
        let upstream = self
            .0
            .select(b"", 256) // hash doesn't matter for round robin
            .unwrap();
        println!("upstream {:?}", upstream);
        let peer = Box::new(HttpPeer::new(upstream, true, "127.0.0.1".to_string()));
        Ok(peer)
    }

    async fn upstream_request_filter(
        &self,
        _session: &mut Session,
        upstream_request: &mut RequestHeader,
        _ctx: &mut Self::CTX,
    ) -> Result<()>
    where
        Self::CTX: Send + Sync,
    {
        upstream_request.insert_header("HOST", "127.0.0.1").unwrap();
        Ok(())
    }
}
